package mylibrary.xml.demo.sample;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.FileInputStream;
import java.security.AlgorithmParameters;
import java.security.Key;
import java.security.spec.AlgorithmParameterSpec;

import javax.crypto.Cipher;
import javax.crypto.SecretKey;
import javax.crypto.SecretKeyFactory;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;
import javax.xml.parsers.DocumentBuilder;
import javax.xml.parsers.DocumentBuilderFactory;
import javax.xml.transform.Transformer;
import javax.xml.transform.TransformerFactory;
import javax.xml.transform.dom.DOMSource;
import javax.xml.transform.stream.StreamResult;

import mylibrary.xml.demo.suit.AlgoNames;
import mylibrary.xml.demo.suit.CipherData;
import mylibrary.xml.demo.suit.EncryptedData;
import mylibrary.xml.demo.suit.EncryptionMethod;
import mylibrary.xml.demo.suit.GenericKeyInfo;

import org.apache.commons.codec.binary.Base64OutputStream;
import org.w3c.dom.Document;
import org.w3c.dom.Node;
import org.w3c.dom.NodeList;

public class XmlEncryption {

	private String algoName = null;
	private String keyName = null;
	private String encId = null;
	private DocumentBuilder docBuilder = null;
	private Document clearDoc = null;
	private Key encKey = null;
	
	private SecretKey decKey = null;
	
	public XmlEncryption() {
		try {
			docBuilder = DocumentBuilderFactory.newInstance().newDocumentBuilder();	
		} catch (Exception e) {
			e.printStackTrace();
		}
	}
	
	private Document getNewDocument() {
		if (docBuilder != null)
			return docBuilder.newDocument();
		else
			return null;	
	}
	
	public void setClearDoc(String fileString) {
		try {
			ByteArrayInputStream in = new ByteArrayInputStream(fileString.getBytes());
			clearDoc = docBuilder.parse(in);	
		} catch (Exception e) {
			e.printStackTrace();
		}
	}
	
	public Document getClearDoc() {
		return this.clearDoc;
	}
	
	public String encryptCompleteXmlFile() {

		// Take an Object of EncryptedData Class.
		// It represents an <EncryptedData> Element.
		EncryptedData encDataObj = this.getEncryptedDataDoc(this.encId, "DOCUMENT");
		// Get XML Structure for <EncryptionMehtod> element.
		Document encMethodDoc = this.getEncryptionMethodDoc(this.algoName);
		// Get XML Structure for <KeyInfo> element.
		Document encKeyInfoDoc = this.getKeyInfoDoc(this.keyName);
		// Read the given file data which will be encrypted.
		String plainData = this.getString(this.clearDoc);
		// Use of JCA/JCE to get encrypted data.
		String cipherData = this.getEncryptedData(plainData);
		// Get XML Structure for <CipherData> element.
		Document cipherDataDoc = this.getCipherDataDoc(cipherData);
		// Join these XML Structures.
		encDataObj.addChild(encMethodDoc);
		encDataObj.addChild(encKeyInfoDoc);
		encDataObj.addChild(cipherDataDoc);
		// return the resultant in a file
		return getString(encDataObj.getEncData());
	}
	
	// Returns the EncryptedData object.
		public EncryptedData getEncryptedDataDoc(String Id, String encType) {
			EncryptedData ed = new EncryptedData(this.getNewDocument());
			ed.setId(Id);
			if (encType.equals("DOCUMENT"))
				ed.setType(AlgoNames.DOCUMENT);
			if (encType.equals("ELEMENT"))
				ed.setType(AlgoNames.ELEMENT);	
			if (encType.equals("CONTENT"))
				ed.setType(AlgoNames.CONTENT);	
			return ed;
		}
		
		
		public Document getEncryptionMethodDoc (String algoName) {
			EncryptionMethod em = new EncryptionMethod(this.getNewDocument());
			if (algoName.equals("DESede"))
				em.setAlgorithm(AlgoNames.TRIPLE_DES);
			return em.getEncMethod();
		}
		
		// Returns the <KeyInfo> structure.
		public Document getKeyInfoDoc (String keyName) {
			GenericKeyInfo ki = 
				new GenericKeyInfo(this.getNewDocument(),"ds", AlgoNames.XML_DSIG);
			ki.setKeyName(keyName);
			return ki.getKeyInfo();	
		}
		
		public Document getCipherDataDoc (String data) {
			CipherData cd = new CipherData(this.getNewDocument());
			cd.setValue(data);
			return cd.getCipherData();
		}
		
		// This is where the actual JCA/JCE data encryption takes place.
		public String getEncryptedData(String data) {
			String clearData = new String(data);
			String cipherDataBase64	= null;
			try {
				//padding the last byte with " "
				int pad = clearData.length()%8;
				for(int i=0;i<(8-pad);i++) 
					clearData += " ";
				// Creating a Cipher Object.
				Cipher cipherObj = Cipher.getInstance(this.algoName+"/CBC/NoPadding");
				// Initializing the Cipher Object in encryption mode.
				cipherObj.init(Cipher.ENCRYPT_MODE, this.encKey);
				// Retreiving the IV
				byte [] iv = cipherObj.getIV();
				//Ciphering the clear text.
				byte[] cipherTemp = cipherObj.doFinal(clearData.getBytes());
				//Prefixing the IV to the cipher string to produce cipherData
				int ivlen = iv.length;
				int cTlen = cipherTemp.length;
				int cDlen = ivlen+cTlen	;
				byte [] cipherData = new byte[cDlen];
				// First 8 bytes are IV bytes
				for (int i=0;i<ivlen;i++) 
					cipherData[i] = (byte)iv[i];
				// Rest of the bytes are cipher bytes	
				for (int i=ivlen;i<cDlen;i++)
					cipherData[i] = (byte)cipherTemp[i-8];	
				//Base 64 encoding of cipherData
				cipherDataBase64 = getBase64Encoded(cipherData);
			}
			catch (Exception e) {
				System.out.println("Problem in getEncryptedData()");
				e.printStackTrace();
			}
			return cipherDataBase64;
		}
		
		//Convert an XmlDocument type object into a String
		public String getString(Document xmldoc) {	
//			String DocumentAsString = null;
//			try	{
//				ByteArrayOutputStream bos = new ByteArrayOutputStream();
//				xmldoc.write(bos);
//				DocumentAsString = bos.toString();
//			}
//			catch(Exception e) {
//				e.printStackTrace();
//			}
//			return DocumentAsString;
			
			ByteArrayOutputStream bos = null;
			try {
				TransformerFactory tf = TransformerFactory.newInstance();
				Transformer t = tf.newTransformer();
				
				bos = new ByteArrayOutputStream();  
	            t.transform(new DOMSource(xmldoc), new StreamResult(bos));
			} catch (Exception e) {
				e.printStackTrace();
			}
            return bos.toString();
		}
		
		private String getBase64Encoded(byte[] plainText) {
			String encoded = null;
			try {
				ByteArrayOutputStream baos = new ByteArrayOutputStream();
				Base64OutputStream out= new Base64OutputStream(baos,true);
				out.write(plainText);
				out.close();
				String encText = new String(baos.toByteArray());
				baos.close();
				encoded = encText;
			}
			catch(Exception e) {
				e.printStackTrace();
			}
			return encoded;
		}
		
		public String getDecryptedData(String encString) {
			String decString = "Under Development";
			try {
				ByteArrayInputStream bais = new ByteArrayInputStream(encString.getBytes());
				Document encDoc = docBuilder.parse(bais);
				
				NodeList nl = encDoc.getElementsByTagName("EncryptedData");
				for(int i=0; i<nl.getLength(); i++) {
					Node edata = nl.item(i);
					String algo = null;
					String keyname = null;
					String encType = null;
					String cipherValue = null;
					
					edata.normalize();
					
					encType = edata.getAttributes().getNamedItem("Type").getNodeValue();
				}
			} catch (Exception e) {
				
			}
		}
		
		public String Decrypt(String encString, Key decKey, String algo) {
			byte[] g = getBase64Decoded(encString);
			int glen = g.length;
			
			byte[] iv = new byte[8];
			for(int i=0; i<9; i++) {
				iv[i] = g[i];
			}
			
			byte[] enc = new byte[glen-8];
			for(int i=0; i<enc.length; i++) {
				enc[i] = g[i+8];
			}
			
			String decString = null;
			try {
				IvParameterSpec ivps = new IvParameterSpec(iv);
				AlgorithmParameters aparam = AlgorithmParameters.getInstance(algo);
				aparam.init(ivps);
				
				Cipher cipher = Cipher.getInstance(algo + "/CBC/NoPadding");
				cipher.init(Cipher.DECRYPT_MODE, decKey, aparam);
				decString = new String(cipher.update(enc));
			} catch (Exception e) {
				e.printStackTrace();
			}
			
		}
		
		public void generateDecKey(String keyName, String algo) {
			try {
				FileInputStream fs = new FileInputStream(keyName);
				byte[] dk = new byte[fs.available()];
				fs.read(dk);
				fs.close();
				
				//SecretKeyFactory keyFactory = SecretKeyFactory.getInstance(algo);
				SecretKey secretKey = new SecretKeySpec(dk, "AES");
				SecretKeyFactory.getInstance("").g
			}
		}
}

